# api/upload.py
from fastapi import APIRouter, Request, File, UploadFile, Form, Depends, HTTPException, status
from nltk.corpus.reader import documents

from api.process_doc import *
from dependencies import *
import os
from llama_index.indices.managed.dashscope import DashScopeCloudIndex

router = APIRouter()

@router.post("/uploadfile")
async def upload_parse_and_create_knowledge_base(file: UploadFile = File(...), knowledge_base_name: str = Form(...)):
    upload_folder = os.getenv("UPLOAD_FOLDER", "uploads")
    os.makedirs(upload_folder, exist_ok=True)
    if file.filename == "":
        raise HTTPException(status_code=400, detail="未选择文件")

    # 校验文件类型
    allowed_extensions = ['.pdf', '.doc', '.docx']
    file_extension = os.path.splitext(file.filename)[1].lower()
    if file_extension not in allowed_extensions:
        raise HTTPException(status_code=400, detail="不支持的文件类型，仅支持 PDF/Doc/Docx 格式")

    # 校验文件大小
    max_file_size = 100 * 1024 * 1024  # 100 MB
    if len(await file.read()) > max_file_size:
        raise HTTPException(status_code=400, detail="文件大小超过100MB")

    # 重置文件指针
    await file.seek(0)
    try:
        file_path = os.path.join(upload_folder, file.filename)
        contents = await file.read()
        with open(file_path, "wb") as buffer:
            buffer.write(contents)
        # 解析文档
        documents = await parse_document(file_path)
        # 处理解析结果
        processed_result = process_parsed_result(documents)
        # 创建知识库索引
        index = DashScopeCloudIndex.from_documents(
            documents,
            knowledge_base_name,
            verbose=True,
        )
        # 查询文档
        dashscope_llm = DashScope(
          model_name=DashScopeGenerationModels.QWEN_MAX, api_key=DASHSCOPE_API_KEY
        )
        query_engine = index.as_query_engine(llm=dashscope_llm)
        abstract = query_engine.query("给出一段300字以内的文本，总结文档内容")
        print(abstract.response)
        return {"parsed_result": processed_result,"abstract": abstract.response}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"文件上传和处理失败，请检查百炼相关配置")

